/*
 * This code implements stVSSIM VQA as suggested by Anush Moorthy and A.C Bovik
 * in "Efficient Motion Weighted Spatio-Temporal Video SSIM Index"
 * http://live.ece.utexas.edu/publications/2010/moorthy_spie_jan10.pdf
 *
 * Please cite the following paper in any published work if you use this software.
 * Anus K. Moorthy, Alan C. Bovik,"Efficient Motion Weighted Spatio-Temporal
 * Video SSIM Index", SPIE Conference on Human Vision and Electronic Imaging,
 * San Jose, CA, January 17-21, 2010
 *
 * File: calcSSIM.cpp
 *
 * This Module calculates the SS-SSIM for an image for the distorted video with
 * reference to the original image
 */

#include "vssim.h"

extern vssim_global_t global;
extern vssim_video_t refVideo;
extern vssim_video_t distortedVideo;
extern vssim_FrScore_t *Sframe;

extern uint32_t findMax(IplImage *img);
extern float calcPSSIM(IplImage *ssim_map);

/*
 * This function calculates the SS-SSIM for the frame passed to it.
 * It calculates the Mu, Sigma and CoVariance between the reference and the
 * distorted frame and returns the final SS-SSIM map for the frame
 */
float calcSSIM(const IplImage *refFr,const IplImage *distortedFr,
				 vssim_size_t size, IplImage *ssim_map)
{
	float C1, C2, K1, K2;
	int L = 256; //dynamic range of the image
	float FrameScore;

	K1 = 0.01;
	K2 = 0.03;
	C1 = (K1*L)*(K1*L);
	C2 = (K2*L)*(K2*L);

	/*convert the 8 bit images to 32bit floats and calculate the image^2*/

	//Convert to 32 bit single channel images
	IplImage *ref = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvConvert(refFr, temp);
	cvResize(temp, ref);
	cvReleaseImage(&temp);

	IplImage *dist = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *tempd = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvConvert(distortedFr, tempd);
	cvResize(tempd, dist);
	cvReleaseImage(&tempd);

	//image square
	IplImage *ref_sq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvPow(ref, ref_sq, 2);

	IplImage *dist_sq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvPow(dist, dist_sq, 2);


	/* Calculate the mean and the variances of the image within 11x11 gaussian window with std deviation = 1.5*/

	int gauss_window = 11;
	double gauss_stddev = 1.5;

	// calc mu as gaussian smooth
	IplImage *ref_mu = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvSmooth(ref, ref_mu, CV_GAUSSIAN, gauss_window, gauss_window, gauss_stddev);

	IplImage *dist_mu = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvSmooth(dist, dist_mu, CV_GAUSSIAN, gauss_window, gauss_window, gauss_stddev);

	// mu sq
	IplImage *ref_musq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvPow(ref_mu, ref_musq, 2);

	IplImage *dist_musq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvPow(dist_mu, dist_musq, 2);


	/* Calculate the sigma of the 11x11 gaussian smoothed window */

	IplImage *ref_sigmasq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvSmooth(ref_sq, ref_sigmasq, CV_GAUSSIAN, gauss_window, gauss_window, gauss_stddev);
	cvAddWeighted(ref_sigmasq, 1, ref_musq, -1, 0, ref_sigmasq);

	IplImage *dist_sigmasq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvSmooth(dist_sq, dist_sigmasq, CV_GAUSSIAN, gauss_window, gauss_window, gauss_stddev);
	cvAddWeighted(dist_sigmasq, 1, dist_musq, -1, 0, dist_sigmasq);


	/* Calculate mu_ref*mu_distorted and covariance between reference and distorted video frames */

	//mu1*mu2
	IplImage *mul_mu = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvMul(ref_mu, dist_mu, mul_mu, 1);

	//sigma12 = cvSmooth(image1*image2)
	IplImage *fr_mul = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvMul(ref, dist, fr_mul, 1);

	IplImage *mul_sigma = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvSmooth(fr_mul, mul_sigma, CV_GAUSSIAN, gauss_window, gauss_window, gauss_stddev);


	/* Calculate the SSIM for Y, U and V separately */
	// Y
	IplImage *ssim = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);

	// (2*mu1*mu2 + C1)
	IplImage *temp1 = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvScale(mul_mu, temp1, 2);
	cvAddS(temp1, cvScalarAll(C1), temp1);

	//(2*sigma12+C2)
	IplImage *temp2 = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvScale(mul_sigma, temp2, 2);
	cvAddS(temp2, cvScalarAll(C2), temp2);

	// (2*mu1*mu2+C1) * (2*sigma12+C2)
	IplImage *temp3 = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvMul(temp1, temp2, temp3, 1);

	// (mu1^2+mu2^2+C1)
	cvAdd(ref_musq, dist_musq, temp1);
	cvAddS(temp1, cvScalarAll(C1), temp1);

	// (sigma1^2+sigma2^2+C2)
	cvAdd(ref_sigmasq, dist_sigmasq, temp2);
	cvAddS(temp2, cvScalar(C2), temp2);

	// (mu1^2+mu2^2+C1) * (sigma1^2+sigma2^2+C2)
	cvMul(temp1, temp2, temp1, 1);

	// (2*mu1*mu2+C1) * (2*sigma12+C2) / (mu1^2+mu2^2+C1) * (sigma1^2+sigma2^2+C2)
	IplImage *temp_ssim = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvDiv(temp3, temp1, temp_ssim, 1);

	// find the max value
	uint32_t maxVal = findMax(temp_ssim);
	cvConvertScale(temp_ssim, ssim, ((float)255/(float)maxVal), 0.0);
	cvResize(ssim, ssim_map);

	FrameScore = calcPSSIM(temp_ssim);

	cvReleaseImage(&temp1);
	cvReleaseImage(&temp2);
	cvReleaseImage(&temp3);
	cvReleaseImage(&ssim);
	cvReleaseImage(&temp_ssim);

	/* clean the memory */

	cvReleaseImage(&ref);
	cvReleaseImage(&ref_mu);
	cvReleaseImage(&ref_musq);
	cvReleaseImage(&ref_sq);
	cvReleaseImage(&ref_sigmasq);

	cvReleaseImage(&dist);
	cvReleaseImage(&dist_mu);
	cvReleaseImage(&dist_musq);
	cvReleaseImage(&dist_sq);
	cvReleaseImage(&dist_sigmasq);

	cvReleaseImage(&fr_mul);
	cvReleaseImage(&mul_mu);
	cvReleaseImage(&mul_sigma);

	return FrameScore;
}

/*
 * This function reads the reference and the distorted files and calls the getSSIM()
 * function which actually calculates the stats and final SS-SSIM score for the frame
 * The same is repeated for U and V frames
 */
vssim_result_t getSSIM(int cntframe)
{
	vssim_size_t size;
	int index = (cntframe/16)-1;

	/* Create the y,u,v images  */
	IplImage *ref_y = cvCreateImage(cvSize(global.size.width, HEIGHT_Y(&global)), IPL_DEPTH_8U, 1);
	IplImage *ref_u = cvCreateImage(cvSize(global.size.width, HEIGHT_U(&global)), IPL_DEPTH_8U, 1);
	IplImage *ref_v = cvCreateImage(cvSize(global.size.width, HEIGHT_V(&global)), IPL_DEPTH_8U, 1);

	IplImage *dist_y = cvCreateImage(cvSize(global.size.width, HEIGHT_Y(&global)), IPL_DEPTH_8U, 1);
	IplImage *dist_u = cvCreateImage(cvSize(global.size.width, HEIGHT_U(&global)), IPL_DEPTH_8U, 1);
	IplImage *dist_v = cvCreateImage(cvSize(global.size.width, HEIGHT_V(&global)), IPL_DEPTH_8U, 1);

	/* Create images for the ssim_map */
	IplImage *y_ssim_map = cvCreateImage(cvSize(global.size.width, HEIGHT_Y(&global)), IPL_DEPTH_8U, 1);
	IplImage *u_ssim_map = cvCreateImage(cvSize(global.size.width, HEIGHT_U(&global)), IPL_DEPTH_8U, 1);
	IplImage *v_ssim_map = cvCreateImage(cvSize(global.size.width, HEIGHT_V(&global)), IPL_DEPTH_8U, 1);

	/* Read the frames into the buffers */
	fseek(refVideo.fptr, 0, SEEK_SET);
	fseek(refVideo.fptr, (cntframe * global.size.width * FRAME_HEIGHT(&global)), SEEK_SET);
	fseek(distortedVideo.fptr, 0, SEEK_SET);
	fseek(distortedVideo.fptr, (cntframe * global.size.width * FRAME_HEIGHT(&global)), SEEK_SET);

	/* read the reference video frame */
	fread(ref_y->imageData, sizeof(char),
			(global.size.width * HEIGHT_Y(&global)),
			refVideo.fptr);

	fread(ref_u->imageData, sizeof(char),
			(global.size.width * HEIGHT_U(&global)),
			refVideo.fptr);

	fread(ref_v->imageData, sizeof(char),
			(global.size.width * HEIGHT_V(&global)),
			refVideo.fptr);

	/* read the distorted video frame*/
	fread(dist_y->imageData, sizeof(char),
			(global.size.width * HEIGHT_Y(&global)),
			distortedVideo.fptr);

	fread(dist_u->imageData, sizeof(char),
			(global.size.width * HEIGHT_U(&global)),
			distortedVideo.fptr);

	fread(dist_v->imageData, sizeof(char),
			(global.size.width * HEIGHT_V(&global)),
			distortedVideo.fptr);

	// Calculate SSIM for Y
	size.width = global.size.width;
	size.height = HEIGHT_Y(&global);
//	printf("\tS-Score - ");
	global.Sframe[index].score_Y = calcSSIM(ref_y, dist_y, size, y_ssim_map);
//	printf("Y: %.3f\t", global.Sframe[index].score_Y);

	// Calculate SSIM for U
	size.width = global.size.width;
	size.height = HEIGHT_U(&global);
	global.Sframe[index].score_U = calcSSIM(ref_u, dist_u, size, u_ssim_map);
//	printf("U: %.3f\t", global.Sframe[index].score_U);

	// Calculate SSIM for V
	size.width = global.size.width;
	size.height = HEIGHT_V(&global);
	global.Sframe[index].score_V = calcSSIM(ref_v, dist_v, size, v_ssim_map);
//	printf("V: %.3f\n", global.Sframe[index].score_V);

	/* Clean up */
	cvReleaseImage(&y_ssim_map);
	cvReleaseImage(&u_ssim_map);
	cvReleaseImage(&v_ssim_map);

	cvReleaseImage(&ref_y);
	cvReleaseImage(&ref_u);
	cvReleaseImage(&ref_v);

	cvReleaseImage(&dist_y);
	cvReleaseImage(&dist_u);
	cvReleaseImage(&dist_v);

	return SUCCESS;
}
